{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction to Entity Linking"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook provides a short tutorial on how to implement and use spaCy's Entity Linking functionality with spaCy v3.  It can be used together with [this video](https://www.youtube.com/watch?v=8u57WSXVpmw). Note that this video was originally created for spaCy v2, but this notebook includes the updates needed for v3. If you want to use spaCy v2 instead, you can find the original code [here](https://github.com/explosion/projects/tree/master/nel-emerson)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Entity Linking** (EL) is the challenge of resolving ambiguous textual mentions to unique concepts in a knowledge base. A related task is **Named Entity Recognition** (NER). An NER component basically identifies words in text that have a specific name and refer to real-world objects, such as people or organizations. spaCy offers pre-built Machine Learning models that perform Named Entity Recognition for a variety of languages (https://spacy.io/models)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install spacy==3.0.6\n",
    "!pip install spacy-lookups-data\n",
    "!python -m spacy download en_core_web_lg"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's load a  pretrained English model, apply it to some sample text and show the named entities that were identified by printing their text and label."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import spacy\n",
    "nlp = spacy.load(\"en_core_web_lg\")\n",
    "text = \"Tennis champion Emerson was expected to win Wimbledon.\"\n",
    "doc = nlp(text)\n",
    "for ent in doc.ents:\n",
    "    print(f\"Named Entity '{ent.text}' with label '{ent.label_}'\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see that this sentence contains a person called \"Emerson\" and an event called \"Wimbledon\". \n",
    "\n",
    "Unfortunately, there may be many people in the world called \"Emerson\", and this output still doesn't tell us which one exactly we meant. This is the challenge addressed by Entity Linking. It transforms an ambiguous textual mention to a unique identifier by looking at the context in which the mention occurs. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this specific case, the sentence gives us important clues: Emerson is clearly a professional tennis player. \n",
    "\n",
    "Searching the internet, we can establish that this sentence is most likely talking about Roy Emerson, an Australian tennis player. We can now resolve this entity in this sentence to its unique identifier from WikiData, which is a free and open, interlingual knowledge base. Its unique IDs always start with a Q, and \"Roy Emerson\" has the identifier Q312545: https://www.wikidata.org/wiki/Q312545"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To implement an entity linking pipeline, we need 3 different steps. \n",
    "\n",
    "The first step, as we already saw, is Named Entity Recognition, in which the mention \"Emerson\" is labeled as a \"Person\". Next, the extracted mention needs to be resolved to a list of plausible candidates. In our case, we'll consider three different people named Emerson. Typically, this list is created by querying a knowledge base (KB) that contains various aliases and synonyms. In the final step, we need to reduce the list of candidates to just one final ID that represents the correct Emerson."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![Diagram of entity linking process](nel_schema.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This tutorial will show you how to use spaCy v3 to create a Knowledge base that will address the second step of candidate generation. Additionally, we will create a new Entity Linking component, and train its Machine Learning model on some annotated data.\n",
    "\n",
    "In this notebook, we implement the functions and training loop from scratch. However, spaCy v3 has introduced a powerful and extensible training configuration system, that we advice to use in most cases. You can find the corresponding implementation with the [config system](https://spacy.io/usage/training#config), runnable with the new [spacy projects](https://spacy.io/usage/projects), [here](https://github.com/explosion/projects/tree/v3/tutorials/nel_emerson).\n",
    "\n",
    "The aim of this tutorial is to help you get started implementing your own Entity Linking functionality with spaCy. If you want to know more about the technical details, checkout this presentation at spaCy IRL 2019: https://www.youtube.com/watch?v=PW3RJM8tDGo&list=PLBmcuObd5An4UC6jvK_-eSl6jCvP1gwXc&index=7&t=0s"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Creating the Knowledge Base "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The first step to perform Entity Linking, is to set up a knowledge base that contains the unique identifiers of the entities we are interested in. In this tutorial we will create a very simple one with only 3 entries. We load the data from a pre-defined CSV file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import csv\n",
    "from pathlib import Path\n",
    "\n",
    "def load_entities():\n",
    "    entities_loc = Path.cwd().parent / \"assets\" / \"entities.csv\"  # distributed alongside this notebook\n",
    "\n",
    "    names = dict()\n",
    "    descriptions = dict()\n",
    "    with entities_loc.open(\"r\", encoding=\"utf8\") as csvfile:\n",
    "        csvreader = csv.reader(csvfile, delimiter=\",\")\n",
    "        for row in csvreader:\n",
    "            qid = row[0]\n",
    "            name = row[1]\n",
    "            desc = row[2]\n",
    "            names[qid] = name\n",
    "            descriptions[qid] = desc\n",
    "    return names, descriptions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "name_dict, desc_dict = load_entities()\n",
    "for QID in name_dict.keys():\n",
    "    print(f\"{QID}, name={name_dict[QID]}, desc={desc_dict[QID]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have 3 entries here, of 3 different people called Emerson. One Australian tennis player, one American writer and one Brazilian footballer. We'll use this information to create our knowledge base. We need to define a fixed dimensionality for the entity vectors, which will be 300-D in our case."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from spacy.kb import KnowledgeBase\n",
    "kb = KnowledgeBase(vocab=nlp.vocab, entity_vector_length=300)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To add each record to the knowledge base, we encode its description using the built-in word vectors of our `nlp` model. The `vector` attribute of a document is the average of its token vectors. We also need to provide a frequency, which is a raw count of how many times a certain entity appears in an annotated corpus. In this tutorial we're not using these frequencies, so we're setting them to an arbitrary value."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for qid, desc in desc_dict.items():\n",
    "    desc_doc = nlp(desc)\n",
    "    desc_enc = desc_doc.vector\n",
    "    kb.add_entity(entity=qid, entity_vector=desc_enc, freq=342)   # 342 is an arbitrary value here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we want to specify aliases or synonyms. We first add the full names. Here, we are 100% certain that they resolve to their corresponding QID, as there is no ambiguity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for qid, name in name_dict.items():\n",
    "    kb.add_alias(alias=name, entities=[qid], probabilities=[1])   # 100% prior probability P(entity|alias)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We also want to add the alias \"Emerson\". We'll assume that each of our 3 Emersons is equally famous and thus we set their probabilities to be equal for each entity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "qids = name_dict.keys()\n",
    "probs = [0.3 for qid in qids]\n",
    "kb.add_alias(alias=\"Emerson\", entities=qids, probabilities=probs)  # sum([probs]) should be <= 1 !"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So this will be our Knowledge base. We can check the entities and aliases that are contained in it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Entities in the KB: {kb.get_entity_strings()}\")\n",
    "print(f\"Aliases in the KB: {kb.get_alias_strings()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also print the candidates that are generated for the full name of Roy Emerson, as well as for the mention \"Emerson\" or for any other random mention, like \"Sofie\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Candidates for 'Roy Stanley Emerson': {[c.entity_ for c in kb.get_alias_candidates('Roy Stanley Emerson')]}\")\n",
    "print(f\"Candidates for 'Emerson': {[c.entity_ for c in kb.get_alias_candidates('Emerson')]}\")\n",
    "print(f\"Candidates for 'Sofie': {[c.entity_ for c in kb.get_alias_candidates('Sofie')]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We notice that querying the KB with the alias \"Emerson\" gives us 3 candidates, but if we query it with an unknown term, it just gives an empty list."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can save the knowledge base by calling the function `to_disk` with an output location."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# change the directory and file names to whatever you like\n",
    "import os\n",
    "output_dir = Path.cwd().parent / \"my_output\"\n",
    "if not os.path.exists(output_dir):\n",
    "    os.mkdir(output_dir) \n",
    "kb.to_disk(output_dir / \"my_kb\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can store the `nlp` object to file by calling `to_disk` as well."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nlp.to_disk(output_dir / \"my_nlp\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Creating a training dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we need to create some annotated data to train an Entity Linking algorithm on. To do so, we will use the annotation tool Prodigy, but you could generate the data in whatever tool you like."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you are watching [the video](https://www.youtube.com/watch?v=8u57WSXVpmw), it will explain how to obtain annotated data with Prodigy. The final result will be a JSONL file that is distributed alongside this notebook. We'll now use this JSONL file to train our entity linker. If you want to skip the annotation part in the video, you can fast forward to [this section](https://www.youtube.com/watch?v=8u57WSXVpmw&t=19m19s)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's have a look at the results in this file:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "from pathlib import Path\n",
    "\n",
    "json_loc = Path.cwd().parent / \"assets\" / \"emerson_annotated_text.jsonl\" # distributed alongside this notebook\n",
    "with json_loc.open(\"r\", encoding=\"utf8\") as jsonfile:\n",
    "    line = jsonfile.readline()\n",
    "    print(line)   # print just the first line"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " We see that the full text of the original sentence is stored, together with a lot of detail about the annotation task. The most important bit is stored with the key `accept` at the end: this is the value of our manual annotation. For this specific sentence and this specific mention, the option with key `Q312545` was manually selected. This is the information that we'll train our entity linker on."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training the Entity Linker"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To feed training data into our Entity Linker, we format our data as a structured tuple. The first part is the raw text, and the second part is a dictionary of annotations. This dictionary defines the named entities we want to link (\"entities\"), as well as the actual gold-standard links (\"links\")."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "from pathlib import Path\n",
    "\n",
    "dataset = []\n",
    "json_loc = Path.cwd().parent / \"assets\" / \"emerson_annotated_text.jsonl\"\n",
    "with json_loc.open(\"r\", encoding=\"utf8\") as jsonfile:\n",
    "    for line in jsonfile:\n",
    "        example = json.loads(line)\n",
    "        text = example[\"text\"]\n",
    "        if example[\"answer\"] == \"accept\":\n",
    "            QID = example[\"accept\"][0]\n",
    "            offset = (example[\"spans\"][0][\"start\"], example[\"spans\"][0][\"end\"])\n",
    "            entity_label = example[\"spans\"][0][\"label\"]\n",
    "            entities = [(offset[0], offset[1], entity_label)]\n",
    "            links_dict = {QID: 1.0}\n",
    "        dataset.append((text, {\"links\": {offset: links_dict}, \"entities\": entities}))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To check whether the conversion looks OK, we can just print the first sample in our dataset. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also check some statistics in this dataset. How many cases of each QID do we have annotated?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gold_ids = []\n",
    "for text, annot in dataset:\n",
    "    for span, links_dict in annot[\"links\"].items():\n",
    "        for link, value in links_dict.items():\n",
    "            if value:\n",
    "                gold_ids.append(link)\n",
    "\n",
    "from collections import Counter\n",
    "print(Counter(gold_ids))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We got exactly 10 annotated sentences for each of our Emersons. Of these, we'll now set aside 6 cases in a separate test set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "\n",
    "train_dataset = []\n",
    "test_dataset = []\n",
    "for QID in qids:\n",
    "    indices = [i for i, j in enumerate(gold_ids) if j == QID]\n",
    "    train_dataset.extend(dataset[index] for index in indices[0:8])  # first 8 in training\n",
    "    test_dataset.extend(dataset[index] for index in indices[8:10])  # last 2 in test\n",
    "    \n",
    "random.shuffle(train_dataset)\n",
    "random.shuffle(test_dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With our datasets now properly set up, we'll now create `Example` objects to feed into the training process. This object is new in spaCy v3. Essentially, it contains a document with predictions (`predicted`) and one with gold-standard annotations (`reference`). During training, the pipeline will compare its predictions to the gold-standard, and update the weights of the neural network accordingly.\n",
    "\n",
    "For entity linking, the algorithm needs access to gold-standard sentences, because the algorithms use the context from the sentence to perform the disambiguation. You can either provide gold-standard `sent_starts` annotations, or run a component such as the `parser` or `sentencizer` on your reference documents:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from spacy.training import Example\n",
    "\n",
    "TRAIN_EXAMPLES = []\n",
    "if \"sentencizer\" not in nlp.pipe_names:\n",
    "    nlp.add_pipe(\"sentencizer\")\n",
    "sentencizer = nlp.get_pipe(\"sentencizer\")\n",
    "for text, annotation in train_dataset:\n",
    "    example = Example.from_dict(nlp.make_doc(text), annotation)\n",
    "    example.reference = sentencizer(example.reference)\n",
    "    TRAIN_EXAMPLES.append(example)\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, we'll create a new Entity Linking component and add it to the pipeline. \n",
    "\n",
    "We also need to make sure the `entity_linker` component is properly initialized. To do this, we need a `get_examples` function that returns some example training data, as well as a `kb_loader` argument. This is a `Callable` function that creates the `KnowledgeBase`, given a certain `Vocab` instance. Here, we will load our KB from disk, using the built-in [`spacy.KBFromFile.v1`](https://spacy.io/api/architectures#KBFromFile) function, which is defined in `spacy.ml.models`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from spacy.ml.models import load_kb\n",
    "\n",
    "entity_linker = nlp.add_pipe(\"entity_linker\", config={\"incl_prior\": False}, last=True)\n",
    "entity_linker.initialize(get_examples=lambda: TRAIN_EXAMPLES, kb_loader=load_kb(output_dir / \"my_kb\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we will run the actual training loop for the new component, taking care to only train the entity linker and not the other components. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from spacy.util import minibatch, compounding\n",
    "\n",
    "with nlp.select_pipes(enable=[\"entity_linker\"]):   # train only the entity_linker\n",
    "    optimizer = nlp.resume_training()\n",
    "    for itn in range(500):   # 500 iterations takes about a minute to train\n",
    "        random.shuffle(TRAIN_EXAMPLES)\n",
    "        batches = minibatch(TRAIN_EXAMPLES, size=compounding(4.0, 32.0, 1.001))  # increasing batch sizes\n",
    "        losses = {}\n",
    "        for batch in batches:\n",
    "            nlp.update(\n",
    "                batch,   \n",
    "                drop=0.2,      # prevent overfitting\n",
    "                losses=losses,\n",
    "                sgd=optimizer,\n",
    "            )\n",
    "        if itn % 50 == 0:\n",
    "            print(itn, \"Losses\", losses)   # print the training loss\n",
    "print(itn, \"Losses\", losses)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The final training loss is pretty small, which is a good sign. But to truly verify whether our model generalizes well, we need to test it on unseen data."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Testing the Entity Linker"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's first apply it on our original sentence. For each entity, we print the text and label as before, but also the disambiguated QID as predicted by our entity linker."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "text = \"Tennis champion Emerson was expected to win Wimbledon.\"\n",
    "doc = nlp(text)\n",
    "for ent in doc.ents:\n",
    "    print(ent.text, ent.label_, ent.kb_id_)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see that Emerson gets disambiguated to Q312545, which is the correct ID for the tennis player. Note also that the entity \"Wimbledon\" gets the annotation `NIL`, which is basically just a placeholder value, showing that the NEL component could not find any relevant ID for this entity. This happens because our Knowledge base and the Entity Linking component have only been trained on \"Emerson\" examples, and are thus quite limited."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's see what the model predicts for the 6 sentences in our test dataset, that were never seen during training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for text, true_annot in test_dataset:\n",
    "    print(text)\n",
    "    print(f\"Gold annotation: {true_annot}\")\n",
    "    doc = nlp(text)  # to make this more efficient, you can use nlp.pipe() just once for all the texts\n",
    "    for ent in doc.ents:\n",
    "        if ent.text == \"Emerson\":\n",
    "            print(f\"Prediction: {ent.text}, {ent.label_}, {ent.kb_id_}\")\n",
    "    print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These results may vary a little from run to run, but usually the EL pipeline will get 5 out of 6 predictions correct (83% accuracy). Random guessing would have only achieved 33%."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Hopefully, this tutorial has shown you how to implement an Entity Linking component in spaCy. The knowledge base and training dataset used here were kept small for demonstration purposes, but in reality you'll want to use a much bigger representative set of entities, perhaps from an ontology or dictionary that is relevant to your use-case. \n",
    "\n",
    "If you have general questions on how to use this functionality in your own application, the best route is to create a new StackOverfow issue and tag it with the label `spaCy`. If you would run into an actual bug with the Entity Linking functionality, you can also open an issue at spaCy's github tracker. \n",
    "\n",
    "I hope your next NLP project will incorporate entity linking !"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.2"
  },
  "pycharm": {
   "stem_cell": {
    "cell_type": "raw",
    "metadata": {
     "collapsed": false
    },
    "source": []
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
